Watch Your Ifs - An InvokeAI Security Vulnerablity Breakdown

This post explains the security vulnerability I discovered and patched inside of the Stable Diffusion-based AI Image Generation tool, InvokeAI.

Background

InvokeAI is a python tool which runs a webserver on a user's computer which is responsible for passing commands to Stable Diffusion via PyTorch. In order to do this, a user must load a large weights file in either SafeTensors or Pickle file formats.

The Problem

Pickle files (.pkl, .ckpt, etc) are extremely unsafe, as they can be trivially crafted to execute arbitrary code when parsed using torch.load

Right now the conventional wisdom among ML researchers and users is to simply not run untrusted pickle files ever, and instead only use SafeTensor files, which cannot (as far as anyone is concerned as of the time of writing) be injected with arbitrary code. This is very good advice.

Unfortunately, I have discovered a vulnerability inside of InvokeAI that allows an attacker to disguise a pickle file as a safetensor and have the payload execute within InvokeAI.

How It Works

Within model_manager.py and convert_ckpt_to_diffusers.py there are if-statements that decide which load method to use based on the file extension of the model file. One path loads the file as a safetensor, and the other loads the file as a pickle. The logic (written in a slightly more readable format than it exists in the codebase) is as follows:

if Path(file).suffix == '.safetensors':
   safetensor_load(file)
else:
   unsafe_pickle_load(file)

Can you spot the issue?

A malicious actor would only need to create an infected pickle file, and then rename the extension to something that does not pass the == '.safetensors' check, but still appears to a user to be a safetensors file so as not to raise suspicion in the user

For example, this might be something like .Safetensors, .SAFETENSORS, SafeTensors, etc.

InvokeAI will look at the file extension, decide it is a Pickle file and not a Safetensor, and then attempt to load it with the Model Manager and then execute the payload.

Proof of Concept

  1. Create a malicious pickle file. Link to example
  2. Rename the .ckpt extension to some variation of .Safetensors, ensuring there is a capital letter anywhere in the extension (eg. malicious_pickle_file.SAFETENSORS)
  3. Import the file like you would normally with any other safetensors file with the Model Manager.
  4. Upon trying to select the model in the web ui, it will be loaded (or attempt to be converted to a Diffuser) with torch.load and the payload will execute.

The Fix

I have merged code into InvokeAI that fixes the problem described above. Now, the model loader will use the safer safetensor_load method by default. Instead of loading as a pickle if the extension is not exactly .safetensors, it will now always load as a safetensors file unless the extension is exactly .ckpt.

Notes:

I think support for pickle files should be totally dropped ASAP as a matter of security, but I understand that there are reasons this would be difficult.

In the meantime, I think the RestrictedUnpickler from Automatic1111, or something similar, should be implemented as a replacement for torch.load, as this significantly reduces the amount of Python methods that an attacker has to work with when crafting malicious payloads inside a pickle file.